import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Rectangle, Polygon

# 示例路径数据（替换为你的实际数据）
# 格式: [x, y, 转角(弧度)]
path = np.array([
    [0.0, 0.0, 0.0],
    [1.0, 0.5, 0.3],
    [2.0, 1.2, 0.6],
    [3.0, 1.0, -0.4],
    [4.0, 0.5, -0.2],
    [5.0, 0.0, 0.1]
])

# 小车参数
CAR_LENGTH = 0.8  # 车长
CAR_WIDTH = 0.4  # 车宽
WHEEL_LEN = 0.3  # 车轮长度
WHEEL_WIDTH = 0.1  # 车轮宽度


def create_car(x, y, theta):
    """创建小车图形组件"""
    # 车身主体（矩形）
    car_body = Rectangle((x - CAR_LENGTH / 2, y - CAR_WIDTH / 2),
                         CAR_LENGTH, CAR_WIDTH,
                         angle=np.degrees(theta), rotation_point='center',
                         color='skyblue', alpha=0.8)

    # 车轮（四个多边形）
    wheels = []
    wheel_positions = [
        (CAR_LENGTH / 2, CAR_WIDTH / 2),  # 右前轮
        (CAR_LENGTH / 2, -CAR_WIDTH / 2),  # 右后轮
        (-CAR_LENGTH / 2, CAR_WIDTH / 2),  # 左前轮
        (-CAR_LENGTH / 2, -CAR_WIDTH / 2)  # 左后轮
    ]

    for wx, wy in wheel_positions:
        # 车轮中心点（局部坐标）
        wheel_center_local = np.array([wx, wy])

        # 旋转到全局坐标系
        rot_matrix = np.array([
            [np.cos(theta), -np.sin(theta)],
            [np.sin(theta), np.cos(theta)]
        ])
        wheel_center_global = rot_matrix @ wheel_center_local + np.array([x, y])

        # 创建车轮（矩形旋转后变成多边形）
        wheel = Rectangle((0, 0), WHEEL_LEN, WHEEL_WIDTH, color='black')
        wheel.set_xy([wheel_center_global[0] - WHEEL_LEN / 2,
                      wheel_center_global[1] - WHEEL_WIDTH / 2])
        wheel.set_angle(np.degrees(theta))
        wheels.append(wheel)

    return car_body, wheels


# 创建图形和坐标轴
fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlim(np.min(path[:, 0]) - 2, np.max(path[:, 0]) + 2)
ax.set_ylim(np.min(path[:, 1]) - 1, np.max(path[:, 1]) + 1)
ax.set_aspect('equal')
ax.grid(True)
ax.set_title('Vehicle Path Tracking Simulation')
ax.set_xlabel('X Position')
ax.set_ylabel('Y Position')

# 绘制路径
ax.plot(path[:, 0], path[:, 1], 'g--', lw=1.5, label='Planned Path')
ax.scatter(path[:, 0], path[:, 1], c='red', s=30, label='Path Points')

# 创建小车初始位置
car_body, wheels = create_car(path[0, 0], path[0, 1], path[0, 2])
ax.add_patch(car_body)
for wheel in wheels:
    ax.add_patch(wheel)

# 添加方向指示器
direction_line, = ax.plot([], [], 'r-', lw=1.5)
direction_arrow = ax.annotate('', xy=(0, 0), xytext=(0, 0),
                              arrowprops=dict(arrowstyle='->', color='red'))


# 动画更新函数
def update(frame):
    """更新每一帧的小车位置和方向"""
    x, y, theta = path[frame]

    # 更新车身位置和方向
    car_body.set_xy([x - CAR_LENGTH / 2, y - CAR_WIDTH / 2])
    car_body.angle = np.degrees(theta)

    # 更新车轮位置
    wheel_positions = [
        (CAR_LENGTH / 2, CAR_WIDTH / 2),
        (CAR_LENGTH / 2, -CAR_WIDTH / 2),
        (-CAR_LENGTH / 2, CAR_WIDTH / 2),
        (-CAR_LENGTH / 2, -CAR_WIDTH / 2)
    ]

    for i, (wx, wy) in enumerate(wheel_positions):
        rot_matrix = np.array([
            [np.cos(theta), -np.sin(theta)],
            [np.sin(theta), np.cos(theta)]
        ])
        wheel_center = rot_matrix @ np.array([wx, wy]) + np.array([x, y])
        wheels[i].set_xy([wheel_center[0] - WHEEL_LEN / 2,
                          wheel_center[1] - WHEEL_WIDTH / 2])
        wheels[i].angle = np.degrees(theta)

    # 更新方向指示器
    arrow_len = CAR_LENGTH * 1.2
    end_x = x + arrow_len * np.cos(theta)
    end_y = y + arrow_len * np.sin(theta)
    direction_line.set_data([x, end_x], [y, end_y])
    direction_arrow.xy = (end_x, end_y)
    direction_arrow.set_position((x, y))

    return [car_body] + wheels + [direction_line, direction_arrow]


# 创建动画
ani = FuncAnimation(fig, update, frames=len(path),
                    interval=200, blit=True)

plt.legend()
plt.tight_layout()
plt.show()

# 如需保存动画（取消注释）
# ani.save('vehicle_simulation.gif', writer='pillow', fps=5)